package com.tangtang.core.tree;

import java.util.ArrayList;
import java.util.Comparator;
import java.util.List;
import java.util.SortedMap;
import java.util.SortedSet;
import java.util.TreeMap;
import java.util.TreeSet;

import com.tangtang.core.basic.Function;

/**
 * Date: 08.09.2011
 */
public class TreeBuilder<T> {

    private List<T> source;
    private final SortedMap<T, SortedSet<T>> childrenMap;
    private final SortedSet<T> roots;

    private final Function<T> getParent;
    private final Comparator<T> sorter;

    public TreeBuilder(List<T> source, Function<T> getParent, Comparator<T> sorter) {
        this.source = source;
        this.getParent = getParent;
        this.sorter = sorter;
        roots = new TreeSet<T>(sorter);
        childrenMap = new TreeMap<T, SortedSet<T>>(sorter);
    }

    public void add(T t) {
        T parent = getParent.apply(source, t);
        if (parent == null) {
            roots.add(t);
        } else {
            SortedSet<T> children = childrenMap.get(parent);
            if (children == null) {
                children = new TreeSet<T>(sorter);
                childrenMap.put(parent, children);
            }
            children.add(t);
            add(parent);
        }
    }

    public List<Tree<T>> build() {
        List<Tree<T>> result = new ArrayList<Tree<T>>();

        for (T root : roots) {
            result.add(new Tree<T>(null, root, childrenMap));
        }
        return result;
    }

    /**
     * 创建builder
     *
     * @param source
     * @param getParent
     * @param sorter
     * @return
     */
    public static <T> TreeBuilder<T> createBuilder(List<T> source, Function<T> getParent, Comparator<T> sorter) {
        TreeBuilder<T> builder = new TreeBuilder<T>(source, getParent, sorter);
        for (T t : source) {
            builder.add(t);
        }
        return builder;
    }

}
